In this implementation, the auto-generation of UUIDs is handled by the PostgreSQL database, specifically within the insert_memory method in the PostgresClient class. When a new record is inserted, the database generates a unique UUID by default, thanks to the id uuid DEFAULT uuid_generate_v4() line in the table creation query. The generated UUID is then returned by the RETURNING id; SQL statement and fetched by return self.cur.fetchone()[0]. This ensures that each memory object has a unique identifier without requiring manual generation.

NOTE: This implementation has many broken tests that will need to be fixed.

The main changes in the create_memory function within the agentmemory/main.py file are as follows:

ID Generation Logic:

The original implementation generated an ID based on the count of documents in the collection and padded it to make it 16 digits long. In the new implementation, this logic has been removed. Instead, the ID is auto-generated by PostgreSQL as a UUID when the record is inserted.

Old:

if id is None:
    id = str(memories.count())
    # pad the id with zeros to make it 16 digits long
    id = id.zfill(16)

New:
The ID generation logic is removed, and the ID (if not provided) is generated by the database.

Insert Operation:

The old implementation used a method called upsert to insert the document into the collection. The new implementation uses a method called add.

Old:

memories.upsert(
    ids=[str(id)],

New:

memories.add(
    ids=ids,
    documents=[text],
    metadatas=[metadata],
    embeddings=[embedding] if embedding is not None else None,
)

Returning ID:

In the old implementation, the function returned the ID generated or provided. In the new implementation, it assumes that the add method appends the newly generated UUID to the ids list, and that is returned.

Old:

return id

New:

memory_id = ids[0]
return memory_id  # This will now be a UUID or the ID you provided

Metadata Handling:

The new version also includes more robust metadata handling, especially for boolean and other non-string types.

By making these changes, the create_memory function now relies on the PostgreSQL database for UUID generation, which is generally more reliable and efficient than manual generation. It also aligns better with industry best practices.

Suggestion for Robust ID Handling

While the new implementation leverages PostgreSQL for UUID generation, it might be beneficial to retain the flexibility of allowing users to specify an ID if they wish. This can be achieved by adding a conditional check before the add method call in the create_memory function.

Here's a suggested modification for create_memory to make it robust to both user-provided IDs and auto-generated UUIDs:

# Check if user-provided ID exists
if id is not None:
    ids = [id]
else:
    ids = []

# Insert the document into the collection
memories.add(
    ids=ids,
    documents=[text],
    metadatas=[metadata],
    embeddings=[embedding] if embedding is not None else None,
)

# If ids list is empty, it means the ID was auto-generated
if not ids:
    memory_id = self.cur.fetchone()[0]  # Fetch the generated UUID from PostgreSQL
else:
    memory_id = ids[0]  # Use the user-provided ID

return memory_id  # This will now be either a UUID or the ID you provided

By implementing this modification, you ensure that if the user provides an ID, that will be used; otherwise, a UUID will be generated by PostgreSQL. This approach preserves the robustness of the original implementation while benefiting from the advantages of database-managed UUIDs.

Resources:

Supabase uuid-ossp extension guide